import cv2, random
import numpy as np
from random import choice

class Indensity(object):
    """
    用于图片增强，随机生成图片增强效果,并随机选取图片增强方法，处理图片
    """
    def __init__(self) -> None:
        self.methods = {"m1": lambda x: self.__flip(x),
                        "m2": lambda x: self.__rotate(x),
                        "m3": lambda x: self.__fuzzy(x),
                        "m4": lambda x: self.__color_trans(x),
                        "m5": lambda x: self.__color_contrast(x),
                        "m6": lambda x: self.__auto_clahe(x),
                        "m7": lambda x: self.__overgray(x),
                        "m8": lambda x: self.__erasure(x),
                        "m9": lambda x: self.__rotation(x),
                        'm10': lambda x: self.__noise(x),
                        'm11': lambda x: self.__convertScaleAbs(x),
                        'm12': lambda x: self.__erode(x),
                        'm13': lambda x: self.__dilate(x),
                        'm14': lambda x: self.__modify_colorChannel(x)}
        self.m_set = list(self.methods)
        self.num = [0, 0, 1, 1, 2, 2, 3, 4, 5]
        self.ways = lambda k: random.sample(self.m_set, k)

    def rise(self, img:np.ndarray)->np.ndarray:
        num = choice(self.num)
        if num == 0:
            return img.copy()
        else:
            for method in self.ways(num):
                img = self.methods[method](img)
            return img


    def __flip(self, img:np.ndarray)->np.ndarray:
        return cv2.flip(img, choice([-1, 0, 1])) # 图片随机按水平、竖直、水平竖直翻转

    def __rotate(self, img:np.ndarray)->np.ndarray:
        # 图片随机旋转90、180、270
        return cv2.rotate(img, choice([cv2.ROTATE_180, 
                                    cv2.ROTATE_90_CLOCKWISE,
                                    cv2.ROTATE_90_COUNTERCLOCKWISE]))

    def __fuzzy(self, img:np.ndarray)->np.ndarray:
        # 图片高斯模糊
        return cv2.GaussianBlur(img, ksize=choice([(3, 3), (5, 5), (7, 7), (9, 9)]),
                                sigmaX=0, borderType=cv2.BORDER_REFLECT)

    def __color_trans(self, img:np.ndarray)->np.ndarray:
        # 颜色变换
        return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    
    def __color_contrast(self, img:np.ndarray)->np.ndarray:
        # 调整对比度
        img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
        h, s, v = cv2.split(img) # 分解hsv图
        v1 = np.clip(v * choice(np.arange(0.45, 2.3, 0.15)),
                    a_min=0,
                    a_max=255) # 随机调整对比
        v1 = v1.astype(np.uint8)
        img = cv2.merge((h, s, v1))
        return cv2.cvtColor(img, cv2.COLOR_HSV2RGB)
    
    def __auto_clahe(self, img:np.ndarray)->np.ndarray:
        # 自适应图像增强，随机增强
        # 动态调整网格大小
        img_h, img_w = img.shape[0], img.shape[1] # 获取图片的shape
        num_h, num_w = max(3, img_h // 50), max(3, img_w // 50)
        b, g, r = cv2.split(img)
        clahe = cv2.createCLAHE(clipLimit=choice([i for i in range(1, 20)]), 
                                tileGridSize=(num_h, num_w))
        b = clahe.apply(b)
        g = clahe.apply(g)
        r = clahe.apply(r)
        return cv2.merge((b, g, r))
    
    def __overgray(self, img:np.ndarray)->np.ndarray:
        img_gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img_gray = cv2.cvtColor(img_gray, cv2.COLOR_GRAY2RGB)
        return cv2.addWeighted(img, alpha=0.65, src2=img_gray, beta=0.35, gamma=0) # 图片叠加
    

    def __erasure(self, img:np.ndarray)->np.ndarray:
        # 图片随机擦除
        # 标记图片起始位置
        x = random.randint(0, int(img.shape[1]*0.95))
        y = random.randint(0, int(img.shape[0]*0.95))
        # 标记长度
        w = random.randint(int(img.shape[1] * 0.05), int(img.shape[1] * 0.35))
        h = random.randint(int(img.shape[0] * 0.05), int(img.shape[0] * 0.35))
        img[y: min(img.shape[0], y+h-1),
            x: min(img.shape[1], x+w-1)] = 0
        return img.copy()
    
    def __noise(self, img:np.ndarray)->np.ndarray:
        # 添加随机噪声
        mean = choice([i for i in range(0, 100)])
        sigma = choice(np.arange(0.1, 2.3, 0.15))
        gaussian = np.random.normal(mean, sigma, img.shape).astype(np.uint8)
        return cv2.add(img, gaussian)
    
    def __rotation(self, img:np.ndarray)->np.ndarray:
        # 图片旋转任意角度（以中心轴）
        angle = np.random.randint(0, 360)
        h, w = img.shape[:2]
        scale_factor = np.random.uniform(1, 1.3)
        m = cv2.getRotationMatrix2D((w/2, h/2), angle, scale_factor)
        return cv2.warpAffine(img, m, (w, h))
    
    def __convertScaleAbs(self, img:np.ndarray)->np.ndarray:
        # 图片随机调整亮度和对比度
        return cv2.convertScaleAbs(img, alpha=choice(np.arange(0.35, 1.7, 0.05)),
                                   beta=choice(np.arange(-70, 70, 5)))

    def __erode(self, img:np.ndarray)->np.ndarray:
        # 图片腐蚀操作
        b, g, r = img[:, :, 0], img[:, :, 1], img[:, :, 2]
        kernel = np.ones((3, 3), np.uint8)
        b = cv2.erode(b, kernel, iterations=1)
        g = cv2.erode(g, kernel, iterations=1)
        r = cv2.erode(r, kernel, iterations=1)
        return cv2.merge([b, g, r])
    
    def __dilate(self, img:np.ndarray)->np.ndarray:
        # 图片腐蚀操作
        b, g, r = img[:, :, 0], img[:, :, 1], img[:, :, 2]
        kernel = np.ones((3, 3), np.uint8)
        b = cv2.dilate(b, kernel, iterations=1)
        g = cv2.dilate(g, kernel, iterations=1)
        r = cv2.dilate(r, kernel, iterations=1)
        return cv2.merge([b, g, r])
    
    def __modify_colorChannel(self, img:np.ndarray)->np.ndarray:
        # 颜色通道随机增加、减少固定数值
        channel = {"b":img[:, :, 0],
                   "g":img[:, :, 1],
                   "r":img[:, :, 2]}
        k = choice([1, 2, 3])
        choose = random.sample(sorted(channel), k) # 随机选择要修改的颜色通道
        for element in choose:
            ch = channel[element].astype(np.int64)
            ch += choice(np.arange(-255, 70, 5))
            ch = np.clip(ch, a_min=0, a_max=255).astype(np.uint8) # 限制元素范围等价于 ch[ch<0] = 0; ch[ch>255] = 255
            channel[element] = ch
        return cv2.merge([channel["b"], channel["g"], channel["r"]])
    
    def __resize(self, img:np.ndarray)->np.ndarray:
        # 随机缩放
        scale_factor = np.random.uniform(0.65, 1.35)
        k = random.choice([0.9, 0.95, 1, 1.05, 1.1])
        return cv2.resize(img, None, fx=scale_factor, fy=scale_factor*k)


    def demo(self, img):
        return self.__modify_colorChannel(img)

if __name__ == "__main__":
    img0 = cv2.imread(r'/home/qhj_work/dataset/indoorCVPR_09/Images/buffet/060926buffet_560.jpg')
    # img0 = cv2.imread(r"C:\Users\QinHJ\Downloads\indoorCVPR_09\videostore\02_02.jpg")
    # img = cv2.GaussianBlur(img, ksize=(11, 11),sigmaX=0, borderType=cv2.BORDER_REFLECT)
    # img = cv2.cvtColor(img0, cv2.COLOR_RGB2BGR)
    # img = cv2.cvtColor(img0, cv2.COLOR_RGB2HLS)
    pro = Indensity()
    # img = pro.rise(img0)
    img = pro.demo(img0)
    cv2.imshow('img', img)
    cv2.imshow('img0', img0)
    cv2.waitKey(0)
    '''
    while True:
        img = pro.rise(img0)
        print(img)'''


